# encoding: utf-8

"""
    计算两段文本的相似度
"""
import numpy as np

from bot.openai_bot import OpenAIBot


def get_similarity(text1: str, text2: str) -> str:
    """
    获取两个短文本的相似度
    :param text1: 文本1
    :param text2: 文本2
    :return:
    """

    openai_bot = OpenAIBot()
    vector1, vector2 = openai_bot.embeddings(text_list=[text1, text2])

    # 计算两个向量的余弦相似度
    cosine_sim = np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))

    return f"{round(cosine_sim * 100, 2)}%"


def main():
    text1 = """The cat jumped over the lazy dog."""
    text2 = """The quick brown fox jumps over the lazy dog."""
    print(get_similarity(text1, text2))

    text1 = """在这些场景下，文本数据是非常丰富和复杂的，使用嵌入向量可以使得计算机更好地理解和处理这些文本数据。"""
    text2 = """通常，当您在应用程序中使用LLM时，您不会将用户输入直接发送到LLM。 相反，您可能正在获取用户输入并构造提示，然后将其发送到LLM。"""
    print(get_similarity(text1, text2))


if __name__ == '__main__':
    main()
